import numpy as np




# 构建keypoint的标准模型输入
def create_inputs(imgs, im_info):
    """generate input for different model type
    Args:
        imgs (list(numpy)): list of image (np.ndarray)
        im_info (list(dict)): list of image info
    Returns:
        inputs (dict): input of model
    """
    inputs = {}
    inputs['image'] = np.stack(imgs, axis=0).astype('float32')
    im_shape = []
    im_pad = []
    for e in im_info:
        im_shape.append(np.array((e['im_shape'])).astype('float32'))
        if "wh_pad" in e.keys():
            pad_ = np.array((e['wh_pad'])).astype('float32')
        else:
            pad_ = np.array((0,0)).astype('float32')
        im_pad.append(pad_)
    inputs['im_shape'] = np.stack(im_shape, axis=0)
    inputs['wh_pad'] = np.stack(im_pad, axis=0)
    return inputs

# paddle-lite 的收入构建函数
def create_inputs_lite(imgs, im_info, batch_size):
    """generate input for different model type
    Args:
        imgs (list(numpy)): list of image (np.ndarray)
        im_info (list(dict)): list of image info
    Returns:
        inputs (dict): input of model
    """
    inputs = {}
    im_c, im_h, im_w = imgs[0].shape[:]
    padding_im = np.zeros((batch_size, im_c, im_h, im_w), dtype=np.float32)
    padding_im[:len(imgs),...] = np.stack(imgs, axis=0).astype('float32')
    inputs['image'] = padding_im
    im_shape = []
    for e in im_info:
        im_shape.append(np.array((e['im_shape'])).astype('float32'))
    inputs['im_shape'] = np.stack(im_shape, axis=0)
    return inputs



def nms(keypoints: np.ndarray, box_radious, match_threshold=0.7):
    """
    对关键点扩充为box然后进行非极大值抑制, 避免关键点重叠
    input:
        keypoints: [n, 3]
    return:
        suppressed: 返回需要抑制的关键点掩码
    """
    # det_num = det_keypoints.shape[0]
    # keypoints = det_keypoints.reshape(-1,3)
    x1 = keypoints[:,0] - box_radious
    y1 = keypoints[:,1] - box_radious
    x2 = keypoints[:,0] + box_radious
    y2 = keypoints[:,1] + box_radious
    scores = keypoints[:,2]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    ndets = keypoints.shape[0]
    suppressed = np.zeros((ndets), dtype=np.int32)

    for _i in range(ndets):
        i = order[_i]
        if suppressed[i] == 1:
            continue
        ix1 = x1[i]
        iy1 = y1[i]
        ix2 = x2[i]
        iy2 = y2[i]
        iarea = areas[i]
        for _j in range(_i + 1, ndets):
            j = order[_j]
            if suppressed[j] == 1:
                continue
            xx1 = max(ix1, x1[j])
            yy1 = max(iy1, y1[j])
            xx2 = min(ix2, x2[j])
            yy2 = min(iy2, y2[j])
            w = max(0.0, xx2 - xx1 + 1)
            h = max(0.0, yy2 - yy1 + 1)
            inter = w * h
            union = iarea + areas[j] - inter
            match_value = inter / union
            if match_value >= match_threshold:
                suppressed[j] = 1
    return suppressed.astype(np.bool8)
